#!/usr/bin/env python3

import argparse
import xarray as xr
import metpy.calc as calc
from metpy.units import units
import numpy as np
import pandas as pd
import pendulum
from scipy.interpolate import interp1d
from glob import glob
import os
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../../utils')
from utils import cli, parse_time, create_kdtree, find_bbox, create_bilinear_wgts, apply_bilinear

def sim_amdar_wrf(files, sid, obs_lon, obs_lat, obs_alt, obs_time, kdtree=None):
	cli.notice(f'Simulate AMDAR.')
	used_files = []
	# Assume the file name contains date time.
	times = []
	for file in sorted(glob(files)):
		f = xr.open_dataset(file, engine='netcdf4')
		time = pendulum.from_format(str(f.Times[0].data)[2:-1], 'YYYY-MM-DD_HH:mm:SS')
		if obs_time[0].subtract(hours=1) <= time <= obs_time[-1].add(hours=1):
			used_files.append(file)
			times.append(time.timestamp())
	f = xr.open_mfdataset(used_files, concat_dim='Time', combine='nested')

	if kdtree == None:
		xlon = f.XLONG[0,:,:]
		xlat = f.XLAT [0,:,:]
		kdtree = create_kdtree(xlon, xlat)

	p = []; u = []; v = []; pt = []; q = []
	for i in range(len(obs_lon)):
		bbox = find_bbox(kdtree, obs_lon[i], obs_lat[i])
		wgts_h = create_bilinear_wgts(xlon, xlat, bbox, obs_lon[i], obs_lat[i])

		# Height levels
		# First, do horizontal interpolation.
		ph      = 0.5 * (f.PH [:,0:f.PH .shape[1]-1,:,:] + f.PH [:,1:f.PH .shape[1],:,:]) / 9.8
		phb     = 0.5 * (f.PHB[:,0:f.PHB.shape[1]-1,:,:] + f.PHB[:,1:f.PHB.shape[1],:,:]) / 9.8
		h       = apply_bilinear(wgts_h, ph + phb, dims=('south_north', 'west_east'))
		tmp1_p  = apply_bilinear(wgts_h, f.P + f.PB, dims=('south_north','west_east')) * units.Pa
		tmp1_u  = apply_bilinear(wgts_h, 0.5 * (f.U[:,:,:,0:f.U.shape[3]-1] + f.U[:,:,:,1:f.U.shape[3]]), dims=('south_north', 'west_east_stag'))
		tmp1_v  = apply_bilinear(wgts_h, 0.5 * (f.V[:,:,0:f.V.shape[2]-1,:] + f.V[:,:,1:f.V.shape[2],:]), dims=('south_north_stag', 'west_east'))
		tmp1_pt = apply_bilinear(wgts_h, f.T + 300, dims=('south_north', 'west_east'))
		tmp1_q  = apply_bilinear(wgts_h, f.QVAPOR, dims=('south_north', 'west_east'))
		# Second, do vertical interpolation.
		tmp2_p = []; tmp2_u = []; tmp2_v = []; tmp2_pt = []; tmp2_q = []
		for j in range(f.Time.size):
			tmp2_p .append(interp1d(h[j,:], tmp1_p [j,:], kind='linear', bounds_error=False, copy=False)(obs_alt[i]))
			tmp2_u .append(interp1d(h[j,:], tmp1_u [j,:], kind='linear', bounds_error=False, copy=False)(obs_alt[i]))
			tmp2_v .append(interp1d(h[j,:], tmp1_v [j,:], kind='linear', bounds_error=False, copy=False)(obs_alt[i]))
			tmp2_pt.append(interp1d(h[j,:], tmp1_pt[j,:], kind='linear', bounds_error=False, copy=False)(obs_alt[i]))
			tmp2_q .append(interp1d(h[j,:], tmp1_q [j,:], kind='linear', bounds_error=False, copy=False)(obs_alt[i]))
		# Third, do time interpolation.
		p .append(interp1d(times, tmp2_p , kind='linear', bounds_error=False, copy=False)(obs_time[i].timestamp()))
		u .append(interp1d(times, tmp2_u , kind='linear', bounds_error=False, copy=False)(obs_time[i].timestamp()))
		v .append(interp1d(times, tmp2_v , kind='linear', bounds_error=False, copy=False)(obs_time[i].timestamp()))
		pt.append(interp1d(times, tmp2_pt, kind='linear', bounds_error=False, copy=False)(obs_time[i].timestamp()))
		q .append(interp1d(times, tmp2_q , kind='linear', bounds_error=False, copy=False)(obs_time[i].timestamp()))
	p  = np.array(p ) * units.Pa
	u  = np.array(u ) * units.mps
	v  = np.array(v ) * units.mps
	pt = np.array(pt) * units.K
	q  = np.array(q )
	# Calculate other quantities.
	ws = calc.wind_speed(u, v)
	wd = calc.wind_direction(u, v)
	t  = calc.temperature_from_potential_temperature(p, pt)
	rh = calc.relative_humidity_from_mixing_ratio(p, t, q).to(units.percent)
	d  = calc.dewpoint_from_relative_humidity(t, rh).to(units.K)
	data = {
		'sid' : sid,
		'lon' : obs_lon,
		'lat' : obs_lat,
		'alt' : obs_alt,
		'time': obs_time,
		'p'   : list(p .data),
		'u'   : list(u .data),
		'v'   : list(v .data),
		't'   : list(t .data),
		'q'   : list(q .data),
		'rh'  : list(rh.data),
		'd'   : list(d .data),
		'ws'  : list(ws.data),
		'wd'  : list(wd.data)
	}
	return pd.DataFrame(data)

def add_obs_err(df, obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p):
	df['u'] = df['u'] + np.random.normal(loc=0, scale=obs_err_u, size=len(df))
	df['v'] = df['v'] + np.random.normal(loc=0, scale=obs_err_v, size=len(df))
	df['t'] = df['t'] + np.random.normal(loc=0, scale=obs_err_t, size=len(df))
	df['p'] = df['p'] + np.random.normal(loc=0, scale=obs_err_p, size=len(df))
	for i in range(len(df)):
		q0 = df['q'][i]
		while True:
			q  = q0 + np.random.normal(loc=0, scale=obs_err_q, size=1)[0]
			if q < 0: continue
			rh = calc.relative_humidity_from_mixing_ratio(df['p'][i] * units.Pa, df['t'][i] * units.K, q)
			if rh > 1: continue
			break
		df.loc[i,['q' ]] = q
		df.loc[i,['rh']] = rh * 100
	df['d' ] = [x.magnitude for x in calc.dewpoint_from_relative_humidity(list(df['t']) * units.K, list(df['rh']) * units.percent).to(units.K)]
	df['ws'] = calc.wind_speed(list(df['u']) * units.mps, list(df['v']) * units.mps)
	df['wd'] = calc.wind_direction(list(df['u']) * units.mps, list(df['v']) * units.mps)
	return df

def write_amdar_littler(file_path, df):
	print(df)
	f = open(file_path, 'w')
	for idx, row in df.iterrows():
		# Header
		f.write('{:20.5f}'.format(row['lat']))                           # latitude
		f.write('{:20.5f}'.format(row['lon']))                           # longitude
		f.write('{:40}'   .format(row['sid']))                           # id
		f.write('{:40}'   .format(row['sid']))                           # name
		f.write('{:40}'   .format('FM-42 AMDAR'))                        # platform
		f.write('{:40}'   .format('OSSE'))                               # source
		f.write('{:>20.5f}'.format(row['alt']))                           # elevation
		f.write('{:>10d}'  .format(1))                                    # num_vld_fld
		f.write('{:>10d}'  .format(0))                                    # num_error
		f.write('{:>10d}'  .format(0))                                    # num_warning
		f.write('{:>10d}'  .format(idx))                                  # seq_num
		f.write('{:>10d}'  .format(0))                                    # num_dups
		f.write('{:>10}'   .format('F'))                                  # is_sound
		f.write('{:>10}'   .format('F'))                                  # bogus
		f.write('{:>10}'   .format('F'))                                  # discard
		f.write('{:>10d}'  .format(-888888))                              # obs_time
		f.write('{:>10d}'  .format(-888888))                              # julian_day
		f.write('{:>20}'   .format(row['time'].strftime('%Y%m%d%H%M%S'))) # date time
		f.write('{:>13.5f}'.format(-888888.0))                            # slp
		f.write('{:>7d}'   .format(0))                                    # slp qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ref_pres
		f.write('{:>7d}'   .format(0))                                    # ref_pres qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ground_t
		f.write('{:>7d}'   .format(0))                                    # ground_t qc
		f.write('{:>13.5f}'.format(-888888.0))                            # sst
		f.write('{:>7d}'   .format(0))                                    # sst qc
		f.write('{:>13.5f}'.format(-888888.0))                            # psfc
		f.write('{:>7d}'   .format(0))                                    # psfc qc
		f.write('{:>13.5f}'.format(-888888.0))                            # precip
		f.write('{:>7d}'   .format(0))                                    # precip qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_max
		f.write('{:>7d}'   .format(0))                                    # t_max qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_min
		f.write('{:>7d}'   .format(0))                                    # t_min qc
		f.write('{:>13.5f}'.format(-888888.0))                            # t_max night
		f.write('{:>7d}'   .format(0))                                    # t_max night qc
		f.write('{:>13.5f}'.format(-888888.0))                            # p_tend03
		f.write('{:>7d}'   .format(0))                                    # p_tend03 qc
		f.write('{:>13.5f}'.format(-888888.0))                            # p_tend24
		f.write('{:>7d}'   .format(0))                                    # p_tend24 qc
		f.write('{:>13.5f}'.format(-888888.0))                            # cloud_cover
		f.write('{:>7d}'   .format(0))                                    # cloud_cover qc
		f.write('{:>13.5f}'.format(-888888.0))                            # ceiling
		f.write('{:>7d}'   .format(0))                                    # ceiling qc
		f.write('\n')
		# Record
		f.write('{:13.5f}'.format(row['p']))                             # pressure (Pa)
		f.write('{:7d}'   .format(0))                                    # pressure qc
		f.write('{:13.5f}'.format(row['alt']))                           # height
		f.write('{:7d}'   .format(0))                                    # height qc
		f.write('{:13.5f}'.format(row['t']))                             # temperature (K)
		f.write('{:7d}'   .format(0))                                    # temperature qc
		f.write('{:13.5f}'.format(row['d']))                             # dewpoint (K)
		f.write('{:7d}'   .format(0))                                    # dewpoint qc
		f.write('{:13.5f}'.format(row['ws']))                            # wind speed (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind speed qc
		f.write('{:13.5f}'.format(row['wd']))                            # wind direction (deg)
		f.write('{:7d}'   .format(0))                                    # wind direction qc
		f.write('{:13.5f}'.format(row['u']))                             # wind u component (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind u component qc
		f.write('{:13.5f}'.format(row['v']))                             # wind v component (m s-1)
		f.write('{:7d}'   .format(0))                                    # wind v component qc
		f.write('{:13.5f}'.format(row['rh']))                            # relative humidity (%)
		f.write('{:7d}'   .format(0))                                    # relative humidity qc
		f.write('{:13.5f}'.format(-888888.0))                            # thickness (m)
		f.write('{:7d}'   .format(0))                                    # thickness qc
		f.write('\n')
		# End
		f.write('{:13.5f}'.format(-777777.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-777777.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('{:13.5f}'.format(-888888.0))
		f.write('{:7d}'   .format(0))
		f.write('\n')
		f.write('{:7d}{:7d}{:7d}\n'.format(1, 0, 0))
	f.close()

def sim_amdar(nature_run_root, nature_run_type,
	sid, lon, lat, alt, time,
	obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p,
	output_littler):
	# Check directory existence.
	if not os.path.isdir(os.path.dirname(output_littler)):
		cli.notice(f'Create directory {os.path.dirname(output_littler)}.')
		os.makedirs(os.path.dirname(output_littler))
	df = sim_amdar_wrf(nature_run_root + '/*', sid, lon, lat, alt, time)
	df = add_obs_err(df, obs_err_u, obs_err_v, obs_err_t, obs_err_q, obs_err_p)
	write_amdar_littler(output_littler, df)

if __name__ == '__main__':
	parser = argparse.ArgumentParser('Simulate AMDAR observation once a time.')
	parser.add_argument('-i', '--nature-run-root', required=True)
	parser.add_argument('-t', '--nature-run-type', choices=('era5', 'wrf'), default='wrf')
	parser.add_argument('-o', '--output-littler', required=True)
	parser.add_argument('--sid', nargs='+')
	parser.add_argument('--lon', nargs='+', type=float)
	parser.add_argument('--lat', nargs='+', type=float)
	parser.add_argument('--alt', help='Flight altitude in meter', nargs='+', type=float)
	parser.add_argument('--time', nargs='+', type=parse_time)
	parser.add_argument('--flight-path-csv')
	parser.add_argument('--obs-err-u', required=True, type=float)
	parser.add_argument('--obs-err-v', required=True, type=float)
	parser.add_argument('--obs-err-t', required=True, type=float)
	parser.add_argument('--obs-err-q', required=True, type=float)
	parser.add_argument('--obs-err-p', required=True, type=float)
	args = parser.parse_args()

	if args.flight_path_csv:
		df = pd.read_csv(args.flight_path_csv, header=None, dtype={0:str,1:float,2:float,3:float,4:str}, parse_dates=[4])
		args.sid = [str(x) for x in df[0]]
		args.lon = df[1]
		args.lat = df[2]
		args.alt = df[3]
		args.time = [pendulum.datetime(x.year, x.month, x.day, x.hour, x.minute, x.second) for x in df[4]]
	elif len(args.sid) != len(args.lon):
		cli.notice(f'Assume there is only one flight {args.sid[0]}.')
		args.sid = [args.sid[0] for _ in args.lon]

	sim_amdar(args.nature_run_root, args.nature_run_type,
		args.sid, args.lon, args.lat, args.alt, args.time,
		args.obs_err_u, args.obs_err_v, args.obs_err_t, args.obs_err_q, args.obs_err_p,
		args.output_littler)
